# -*- coding: utf-8 -*-
# coding=utf-8
import os
import rasterio
import geopandas as gpd
from rasterio.mask import mask
from rasterio.plot import show
from rasterio.windows import Window, from_bounds
from shapely.geometry import Point, Polygon, box, mapping, LineString
import numpy as np

'''
这里的矩形可以自己选择大小，当然也可以通过读取矢量的空间范围去裁剪。同时代码中还加入了压缩参数。
'''
def compress_tif(tif_path, output_file):
    with rasterio.open(tif_path) as src:
        transform = src.transform
        # minx, miny, maxx, maxy = geometry1.bounds(要素的地理字段)
        expanded_bbox = box(minx=0, miny=0, maxx=100, maxy=100)  # 创建一个新的扩展后的边界框（Shapely box）
        window = from_bounds(expanded_bbox.bounds[0], expanded_bbox.bounds[1],
                             expanded_bbox.bounds[2], expanded_bbox.bounds[3],
                             transform=transform)  # 将Shapely几何对象转换为rasterio可以理解的边界
        clipped = src.read(window=window)  # 读取并裁剪TIFF数据

        out_meta = src.meta.copy()
        out_meta.update({"driver": "GTiff", "compress": 'lzw'})  # 更新元数据,rle，lzw等
        del src
    with rasterio.open(output_file, 'w', **out_meta) as dest:  # 写入裁剪后的TIFF文件
        dest.write(clipped)


# 压缩和裁剪栅格数据
def compress_tif(tif_path, output_file):
    with rasterio.open(tif_path) as src:
        transform = src.transform
        # minx, miny, maxx, maxy = geometry1.bounds(要素的地理字段)
        # expanded_bbox = box(minx=0, miny=0, maxx=100, maxy=100)  # 创建一个新的扩展后的边界框（Shapely box）
        # window = from_bounds(expanded_bbox.bounds[0], expanded_bbox.bounds[1],
        #                      expanded_bbox.bounds[2], expanded_bbox.bounds[3],
        #                      transform=transform)  # 将Shapely几何对象转换为rasterio可以理解的边界
        window = from_bounds(11, 2, 3, 4, transform=transform)
        clipped = src.read(window=window)  # 读取并裁剪TIFF数据

        out_meta = src.meta.copy()
        out_meta.update({"driver": "GTiff", "height": window.height, "width": window.width,
                         "transform": rasterio.windows.transform(window, transform), "compress": 'lzw'})
        # 更新元数据,rle，lzw等
        del src
    with rasterio.open(output_file, 'w', **out_meta) as dest:  # 写入裁剪后的TIFF文件
        dest.write(clipped)
'''
   这里使用到了geopandas库用来读取每个要素的空间范围。
'''
def clip_raster_from_features(vector_file=r"彭俊喜/1.shp", raster_file=r"彭俊喜/1.tif"):
    """
    :param vector_file: 输入需要裁剪的面矢量（多面）
    :param raster_file: 输入需要裁剪的影像
    :return: None
    """
    # 读取面矢量数据
    gdf = gpd.read_file(vector_file)
    # 循环遍历面矢量中的每个面要素
    for index, row in gdf.iterrows():
        # 获取当前面要素的几何形状
        geometry1 = row.geometry
        # 确保面要素不是空的
        if geometry1.is_empty:
            print("Skipping empty geometry for feature %d" %(index))
            continue
            # 打开影像文件
        with rasterio.open(raster_file) as src:
            # 将面要素的边界转换为shapely的box对象
            geometry_bounds = box(*geometry1.bounds)
            # 将rasterio的bounds转换为shapely的box对象
            src_bounds = box(*src.bounds)
            # 检查面要素的边界是否与影像的边界相交
            if not geometry_bounds.intersects(src_bounds):
                print("Skipping feature {index} as it does not intersect with the raster.")
                continue
                # 转换几何形状为Rasterio可以理解的格式
            geom_for_rasterio = mapping(geometry1)
            # 使用面要素裁剪影像
            try:
                out_image, out_transform = mask(src, [geom_for_rasterio], crop=True)
                out_image[out_image == src.nodata] = np.nan
            except ValueError as e:
                print("Error clipping feature %d: %s"%(index,e))
                continue
                # 检查是否成功裁剪出影像
            if out_image is None:
                print("No data was clipped for feature %d. Skipping."%(index))
                continue
                # 为裁剪后的影像设置输出路径和文件名
            output_file = 'clipped_image_%d.tif'%(index)
            output_path = os.path.join('G:\\code\\ydh\\py_tiff_back\\testData\\tiffData\\',
                                       output_file)
            # 创建输出文件，并写入裁剪后的影像数据
            dtypes = [src.dtypes[i] for i in range(src.count)]  # 确保数据类型列表与波段数量匹配
            # 假设所有波段的数据类型都是相同的，并且你想要保持与输入影像相同的数据类型
            dtype = src.dtypes[0]  # 获取第一个波段的数据类型
            # 使用这个数据类型打开输出文件
            with rasterio.open(output_path, 'w', driver='GTiff', height=out_image.shape[1],
                               width=out_image.shape[2], count=src.count, dtype=dtype,
                               crs=src.crs, transform=out_transform) as dest:
                dest.write(out_image)
            print('Clipped image %s has been saved.'%(output_file))
    print('All features have been processed.')
tiffPath = "G:\\code\\ydh\\py_tiff_back\\testData\\tiffData\\RS_FOR_BJ_FMP_202409030610_202409031110_1H_5_1_1.tif"
beijingShpPath = "G:\\code\\ydh\\py_tiff_back\\testData\\shp\\beijing4490.shp"
clip_raster_from_features(beijingShpPath,tiffPath)